package com.ng8866.gateway.controller;

import cn.hutool.json.JSONObject;
import com.ng8866.gateway.balancer.LeastConnectionLoadBalancer;
import org.springframework.data.repository.query.Param;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

@RestController
@RequestMapping("/javastats")
public class StatsController {
    /**
     * 加载服务映射表
     *
     * @param serviceName 服务名称，用于获取特定服务的映射表
     * @return 包含服务映射表的JSONObject对象，若未找到对应服务或loadMap为空则返回空的JSONObject对象
     */
    @RequestMapping("/loadMap")
    public JSONObject loadMap(@Param("serviceName") String serviceName) {
        JSONObject ret = new JSONObject();
        Map<String, ConcurrentHashMap<String, AtomicInteger>> loadMap = LeastConnectionLoadBalancer.loadMap;
        if (loadMap != null) {
            if (serviceName != null && !serviceName.isEmpty()) {
                ConcurrentHashMap<String, AtomicInteger> serviceMap = loadMap.get(serviceName);
                if (serviceMap != null) {
                    ret = new JSONObject(serviceMap);
                } else {
                    ret = new JSONObject();
                }
            } else {
                // 如果需要返回整个loadMap
                ret = new JSONObject(loadMap);
            }
        }
        return ret;
    }

    /**
     * 获取负载均衡器中的连接计数映射表，并转换为JSONObject返回
     *
     * @return 包含连接计数映射表的JSONObject对象，如果映射表为null则返回空的JSONObject对象
     */
    @RequestMapping("/countMap")
    public JSONObject countMap() {
        ConcurrentHashMap<String, AtomicInteger> countMap = LeastConnectionLoadBalancer.countMap;
        if (countMap == null) {
            return new JSONObject(); // 返回空的JSONObject当countMap为null
        }
        // 转换countMap为JSONObject并返回
        return new JSONObject(countMap);
    }

    /**
     * 设置连接负载均衡器的计数状态
     *
     * @param count 计数状态，true 表示启用计数，false 表示禁用计数
     * @return 返回一个空的 JSONObject 对象
     */
    @RequestMapping("/setCountMap")
    public JSONObject setCountMap(@Param("count") Boolean count) {
        LeastConnectionLoadBalancer.setCount(count);
        return (new JSONObject()).set("count", LeastConnectionLoadBalancer.isCount);
    }

}
